package leetcode; /**
 * description:
 * author:张腾
 * date:2021-07-28
 */

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 给定一个二叉树（具有根结点root），一个目标结点target，和一个整数值 K 。
 *
 * 返回到目标结点 target 距离为 K 的所有结点的值的列表。 答案可以以任何顺序返回。
 */
public class leetcode863 {
    Map<Integer,TreeNode> parents = new HashMap<>();
    List<Integer> res = new ArrayList<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        findParents(root);
        search(target,null,0,k);
        return res;
    }

    public void findParents(TreeNode root){
        if (root.left!=null){
            parents.put(root.left.val,root);
            findParents(root.left);
        }
        if (root.right!=null){
            parents.put(root.right.val,root);
            findParents(root.right);
        }
    }

    public void search(TreeNode root,TreeNode from,int depth,int k){
        if (root==null) return;
        if (depth==k) {
            res.add(root.val);
            return;
        }
        if (root.left!=from)
            search(root.left,root,depth+1,k);
        if (root.right!=from)
            search(root.right,root,depth+1,k);
        if (parents.get(root)!=from){
            search(parents.get(root.val),root,depth+1,k);
        }
    }

}
